﻿using System;
using System.Threading;
using System.Threading.Tasks;

namespace CuteAnt.AsyncEx
{
  /// <summary>Holds the task for a cancellation token, as well as the token registration. The registration is disposed when this instance is disposed.</summary>
  public sealed class CancellationTokenTaskSource<T> : IDisposable
  {
    /// <summary>The cancellation token registration, if any. This is <c>null</c> if the registration was not necessary.</summary>
    private readonly IDisposable _registration;

    /// <summary>Creates a task for the specified cancellation token, registering with the token if necessary.</summary>
    /// <param name="cancellationToken">The cancellation token to observe.</param>
    public CancellationTokenTaskSource(CancellationToken cancellationToken)
    {
      if (cancellationToken.IsCancellationRequested)
      {
#if NET_4_5_GREATER
        Task = System.Threading.Tasks.Task.FromCanceled<T>(cancellationToken);
#else
        Task = TaskConstants<T>.Canceled;
#endif
        return;
      }
      var tcs = new TaskCompletionSource<T>();
#if NET_4_5_GREATER
      _registration = cancellationToken.Register(() => tcs.TrySetCanceled(cancellationToken), useSynchronizationContext: false);
#else
      _registration = cancellationToken.Register(() => tcs.TrySetCanceled(), useSynchronizationContext: false);
#endif
      Task = tcs.Task;
    }

    /// <summary>Gets the task for the source cancellation token.</summary>
    public Task<T> Task { get; private set; }

    /// <summary>Disposes the cancellation token registration, if any. Note that this may cause <see cref="Task"/> to never complete.</summary>
    public void Dispose()
    {
      _registration?.Dispose();
    }
  }
}
